#define BLOCK_SIZE 128

#include <iostream>
#include <vector>

#include "radixSort.h"

__global__ void scan(int* data, int* totals, int n)
{
	//extern __shared__ int temp[];
	__shared__ int temp[BLOCK_SIZE];

	int block_offset = blockIdx.x*blockDim.x;
	int t_id = threadIdx.x;
	int x = block_offset + t_id;

	// if n is not a power of 2, pad local memory with 0
	if (x < n)
		temp[t_id] = data[x];
	else
		temp[t_id] = 0;

	int offset = 1;

	for (int d = blockDim.x >> 1; d > 0; d >>= 1)
	{
		__syncthreads();

		if (t_id < d)
		{
			int a = offset*(2*t_id + 1) - 1;
			int b = a + offset;

			temp[b] += temp[a];
		}

		offset *= 2;
	}

	if (t_id == 0)
	{
		totals[blockIdx.x] = temp[blockDim.x - 1];
		temp[blockDim.x - 1] = 0;
	}

	for (int d = 1; d < blockDim.x; d *= 2)
	{
		offset >>= 1;
		__syncthreads();

		if (t_id < d)
		{
			int a = offset*(2*t_id + 1) - 1;
			int b = a + offset;

			int t = temp[a];
			temp[a] = temp[b];
			temp[b] += t;
		}
	}

	__syncthreads();

	if (x < n)
		data[x] = temp[t_id];
}

__global__ void addTotals(int* data, int* totals, int n)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;

	if (x < n)
		data[x] += totals[blockIdx.x];
}

int prefixSum(int* d_data, int n)
{
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((n-1)/dimb.x + 1, 1);

	int* d_totals;
	cudaMalloc(&d_totals, dimg.x*sizeof(int));

	scan<<<dimg, dimb>>>(d_data, d_totals, n);

	int total = 0;

	if (dimg.x > 1)
	{
		total = prefixSum(d_totals, dimg.x);
		addTotals<<<dimg, dimb>>>(d_data, d_totals, n);
	}
	else
	{
		cudaMemcpy(&total, d_totals, sizeof(int), cudaMemcpyDeviceToHost);
	}

	cudaFree(d_totals);

	return total;
}



__global__ void count(record_type* records, int* counts, int n, int mask)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;

	if (x < n)
	{
		if ((records[x].first & mask) == 0)
			counts[x] = 1;
		else
			counts[x] = 0;
	}
}

__global__ void split(record_type* records, record_type* splitted, int* offsets, int total, int n, int mask)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;

	if (x < n)
	{
		if ((records[x].first & mask) == 0)
			splitted[offsets[x]] = records[x];
		else
			splitted[x - offsets[x] + total] = records[x];
	}
}

void radixSort(record_type* d_records, int n)
{
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((n-1)/dimb.x + 1, 1);

	int* d_offsets;
	record_type* d_temp;
	cudaMalloc(&d_offsets, n*sizeof(int));
	cudaMalloc(&d_temp, n*sizeof(record_type));

	unsigned int mask = 1;
	for (int i = 0; i != 32; ++i, mask <<= 1)
	{
		count<<<dimg, dimb>>>(d_records, d_offsets, n, mask);
		int total = prefixSum(d_offsets, n);
		split<<<dimg, dimb>>>(d_records, d_temp, d_offsets, total, n, mask);

		std::swap(d_records, d_temp);
	}

	cudaFree(d_offsets);
	cudaFree(d_temp);
}
